# %%
import numpy as np
from matplotlib import pyplot as plt
import WeylPoint_Count
import socket
import time
import scipy.io as sio
import socket
import os
import Diag_Hamiltonian_No_psi
from variables import (ku_grid, ky_grid, kv_grid, Nu, Ny, Nv, m0, n0, l0, version_count, gridN, 
                       save_address, num_of_band, V0, Vy_List, Omega0, mz_list, psi_list, AtomTemperture,
                       AtomDensityHot)
# %%
time_start = time.time()
time_start_struct = time.localtime(time_start)
# %%
os.system('chcp 65001')
Host_Name = socket.gethostname()
if Host_Name == 'VictorChengPC':
    Computer_Name = 'X:\\Work\\mail.ustc.edu.cn\\Rb87BEC - Files\\'
elif Host_Name == 'calcenter2':
    Computer_Name = 'D:\\vc\\mail.ustc.edu.cn\\Rb87BEC - Files\\'
elif Host_Name == 'Rb87BEC3':
    Computer_Name = '/media/rb87bec3/Files/vc/Rb87PythonData/3DSOCbandData/'
elif Host_Name == 'ip-172-31-42-28':
    Computer_Name = '/home/ubuntu/Data/'
Threshold = 0.025

# %%
energies_All = np.zeros((2, 2, gridN, 2, len(Vy_List)*len(mz_list)))
Number_Map = np.zeros((len(Vy_List), len(mz_list)))
for Vymz_Index in np.flip(range(len(Vy_List)*len(mz_list))):
    Indices = np.unravel_index(Vymz_Index, Number_Map.shape)
    Vy_Index = Indices[0]
    mz_Index = Indices[1]
    Vy = Vy_List[Vy_Index]
    mz = mz_list[mz_Index]
    del_psi = psi_list[0]
    Vxz = V0 * np.cos(del_psi)
    Omega_xy = Omega0 * np.sin(del_psi / 2 + np.pi / 4)
    Omega_zy = Omega0 * np.cos(del_psi / 2 + np.pi / 4)
    t1 = Vxz
    t2 = Vy
    t3 = Omega_xy
    t4 = Omega_zy
    if Host_Name == 'VictorChengPC' or Host_Name == 'calcenter2':
        Load_Name = Computer_Name + "Rb87Data\\3DSOCbandData\\190926\\3x3x3;4x3x%dx3;" \
            "V_xz=1.770E_r,V_y=%.3fE_r,Oxy=1.022E_r,Ozy=-1.022E_r,T=1.000e-07K,n=3.000e+18m-3" \
            "del_psi=1.000pi,mz=%.3fErno_states_v1.npz" % (gridN, Vy, mz)
    elif Host_Name == 'Rb87BEC3':
        Load_Name = Computer_Name + "/190926/3x3x3;4x3x%dx3;" \
            "V_xz=1.770E_r,V_y=%.3fE_r,Oxy=1.022E_r,Ozy=-1.022E_r,T=1.000e-07K,n=3.000e+18m-3" \
            "del_psi=1.000pi,mz=%.3fErno_states_v1.npz" % (gridN, Vy, mz)
    if os.path.isfile(Load_Name) and 0:
        Data = np.load(Load_Name)
        ku_grid = Data['ku_grid']
        ky_grid = Data['ky_grid']
        kv_grid = Data['kv_grid']
        energies = Data['energies']
        energies = energies[0:2, ...]
    else:
        energies, ESigmaZ = Diag_Hamiltonian_No_psi.diag_no_states(mz, num_of_band, t1, t2, t3, t4)
    energies_All[..., Vymz_Index] = energies
    gap = energies[1, ...] - energies[0, ...]

    Weyl_Points = []
    Number_of_WeylPoints = 0
    for ku_index in range(len(ku_grid)):
        if ku_grid[ku_index] == 1:
            continue
        gap2d = gap[ku_index, ...]
        Weyl_Points_ku, Number_of_WeylPoints_ku = WeylPoint_Count.Count(gap2d, ky_grid, kv_grid, Threshold)
        Weyl_Points_kup = []
        if Number_of_WeylPoints_ku != 0:
            for point in Weyl_Points_ku:
                Weyl_Points_kup.append((ku_grid[ku_index], point[0], point[1]))
        Number_of_WeylPoints += Number_of_WeylPoints_ku
        Weyl_Points.extend(Weyl_Points_kup)
    # Number_of_WeylPoints = Data['Number_of_WeylPoints']
    Number_Map[Vy_Index, mz_Index] = Number_of_WeylPoints
    if np.mod(Vymz_Index, 100) == 0:
        print('%d out of %d' % (Vymz_Index, len(Vy_List)*len(mz_list)))
fig = plt.imshow(Number_Map, cmap='seismic')
plt.colorbar()
plt.show()
os.mkdir()
if Host_Name == 'VictorChengPC' or Host_Name == 'calcenter2':
    sio.savemat(Computer_Name + r"Rb87Data\3DSOCbandData\190926\Number_Map,gridN=%d,"
                r"Threshold=%.3f.mat" % (gridN, Threshold),
                {'Number_Map': Number_Map})
elif Host_Name == 'rb87bec3':
    pass
elif Host_Name == 'ip-172-31-42-28':
    sio.savemat(Computer_Name + r"Number_Map,gridN=%d,"
                                r"Threshold=%.3f.mat" % (gridN, Threshold),
                {'Number_Map': Number_Map, 'energies_All': energies_All})

# %%
time_end = time.time()
time_end_struct = time.localtime(time_end)
print('Time spent: ', time_end - time_start, 's')
print('Start time:', time_start_struct)
print('End time:', time_end_struct)
